#include "clamav.h"
#include "matcher-bloom.h"

// optimum values for CPUs with 256 KB cache sizes
const int BETA = 7;
const int BLOOM_FILTER_SIZE = 4*1024*1024;

// 1.5 * 4 =~ 6 MB memory overhead
const int HASHTABLE_PRIME = 1572869;

bit *ba_new(const elem_t nelems) 
{
  size_t howmany =  NELEM(nelems,(BITS_SZ));
   
  return ((bit *)calloc(howmany, sizeof(bit)));
} 

void ba_assign(bit arr[], elem_t elem, const boolean value) 
{
  if (value) 
    arr[elem / BITS_SZ] |= (1 << (elem % BITS_SZ));
  else 
    arr[elem / BITS_SZ] &= ~(1 << (elem % BITS_SZ));
}

inline
boolean ba_value(const bit arr[], const elem_t elem) 
{
  return( (arr[elem / BITS_SZ] & (1 << (elem % BITS_SZ))) ? TRUE:FALSE );
} 

void init_bloom_filter(bit *to_init)
{
  int i;
  for (i = 0; i < BLOOM_FILTER_SIZE; i++)
    ba_assign(to_init, i, FALSE);
}


// 4 very fast hash functions: Mask (literally inlined), xor, fasthash and
// sdbm. 

// len can NOT be shorter than 2+4 = 6 bytes. This puts a lower limit on how
// small Beta can be. There should be a check in this function to stop it from
// getting called with len < 6, but it's omitted for performance reasons (an
// extra branch instruction is very costly here).
inline
unsigned int xor_hash(unsigned char *buf, size_t len)
{
  unsigned int hash_val = 0;

  hash_val ^= *((unsigned int *) (buf));
  hash_val ^= *((unsigned int *) (buf+1));
  hash_val ^= *((unsigned int *) (buf+2));

  return hash_val;
}

// fast hash from hashlib.c
inline
unsigned long fast_hash(unsigned char *buf, size_t len)
{
  unsigned long hval = 0;
 
  unsigned char *bptr = buf;	/* start of buffer */
  unsigned char *bend = bptr + len;		/* beyond end of buffer */
  
  while (bptr < bend)
    hval = hval * 31UL + *bptr++;
 
  return hval;
}

// sdbm
inline
unsigned long sdbm(unsigned char *buf, size_t len)
{
  unsigned long hval = 0;
  int c;

  unsigned char *bptr = buf;	/* start of buffer */
  unsigned char *bend = bptr + len;		/* beyond end of buffer */
  
  while (bptr < bend)
  {
    c = *bptr++;
    hval = c + (hval << 6) + (hval << 16) - hval;
  }

  return hval;
}

boolean is_possible_match(bit *filter_to_search, unsigned char* string_to_match)
{  
  int mask1024 = 0x3FFFFF;

  if (ba_value(filter_to_search, (fast_hash(string_to_match, BETA) & mask1024))
      == FALSE)
    return FALSE;

  if (ba_value(filter_to_search, (sdbm(string_to_match, BETA) & mask1024))
      == FALSE)
    return FALSE;

  return TRUE;
}

// Inserts entries into the bloom-filter
void set_bloom_filter_value(bit *filter_to_set, unsigned char *string_to_hash)
{
  int mask1024 = 0x3FFFFF;

  ba_assign(filter_to_set,
	    ((*(unsigned int *) string_to_hash) & mask1024), TRUE);
  ba_assign(filter_to_set, (xor_hash(string_to_hash, BETA) & mask1024), TRUE);
  ba_assign(filter_to_set, (fast_hash(string_to_hash, BETA) & mask1024), TRUE);
  ba_assign(filter_to_set, (sdbm(string_to_hash, BETA) & mask1024), TRUE);
}

void insert_into_hash_table(struct cli_bloom_node* hash_table[],
			    const char *virname, const char *hexsig, 
			    unsigned short type)
{
  if ((strlen(hexsig) / 2) < BETA)
  {
    cli_errmsg("BloomAV can't scan for very short virus sigs! Skipping.");
    return;
  }

  struct cli_bloom_node* bloom_node = 
    (struct cli_bloom_node*) malloc(sizeof(struct cli_bloom_node));

  bloom_node->signature_key = cli_hex2str((char*) hexsig);
  bloom_node->signature_length = strlen(hexsig) / 2;
  bloom_node->signature_name = strdup(virname);
  bloom_node->next = NULL;

  int hash_key = sdbm(bloom_node->signature_key, BETA) % HASHTABLE_PRIME;

  // if item was already in the hash-table, do a list head insertion
  if ( hash_table[hash_key] )
  {
    struct cli_bloom_node* second_node = hash_table[hash_key]->next;
    hash_table[hash_key]->next = bloom_node;
    bloom_node->next = second_node;
  }
  else
  {
    hash_table[hash_key] = bloom_node;
  }

  // Error checking for memory allocations omitted for now.
  // return 0;
}

// returns 1 if virus is found
// 0 otherwise
int lookup_in_hash_table(const char* rest_of_file, int remaining_size,
			 const char **virname, 
			 struct cli_bloom_node* hash_table[])
{
  int hash_key = sdbm((unsigned char*) rest_of_file, BETA) % HASHTABLE_PRIME;

  struct cli_bloom_node* table_iter = hash_table[hash_key];
  while (table_iter)
  {
    if (remaining_size < table_iter->signature_length)
    {
      table_iter = table_iter->next;
      continue;
    }

    if (!memcmp(rest_of_file, table_iter->signature_key,
		table_iter->signature_length))
    {
      (*virname) = (char*) table_iter->signature_name;
      return 1;
    }
 
    table_iter = table_iter->next;
  }

  // else there is a false positive from Bloom filter, gather stats
  // else { }

  return 0;
}

void setup_bloom_structures(bit** bloom_filter, 
			    struct cli_bloom_node*** hash_table)
{
  // Add Error checking into this function.

  (*bloom_filter) = ba_new(BLOOM_FILTER_SIZE);
  init_bloom_filter(*bloom_filter);

  (*hash_table) = (struct cli_bloom_node**) 
    calloc(HASHTABLE_PRIME, sizeof(struct regular_signature*));
}

void clean_bloom_structures(bit* bloom_filter, 
			    struct cli_bloom_node* hash_table[]) 
{
  free(bloom_filter);

  int i = 0;
  for (i = 0; i < HASHTABLE_PRIME; i++)
  {
    if (hash_table[i])
    {
      struct cli_bloom_node* list_head = hash_table[i];

      while (list_head)
      {
	struct cli_bloom_node* sig_to_delete = list_head;
	list_head = list_head->next;

	free(sig_to_delete->signature_key);
	free(sig_to_delete->signature_name);
	
	free(sig_to_delete);
      }
    }
  }

  free(hash_table);
}

int cli_bloom_filter_scanbuff(const char *buffer, unsigned int length, 
			      const char **virname, const struct cl_node *root,
			      unsigned short ftype)
{
  bit* bloom_filter = root->bloom_filter;
  struct cli_bloom_node** sig_hashtable = root->sig_hashtable;

  // this is where we do the scanning.
  char *file_iterator = (char *) buffer;
  long int file_size_ctr = 0;

  while (file_size_ctr < length)
  {

    // The fastest hash function is literally inlined.
    if ((ba_value(bloom_filter, ((*(unsigned int *) file_iterator) &
				 0x3FFFFF)) == TRUE) &&
	(ba_value(bloom_filter,(xor_hash(file_iterator, BETA) &
				0x3FFFFF)) == TRUE))
    {
      if (is_possible_match(bloom_filter, file_iterator) == TRUE)
      {
	if (lookup_in_hash_table((const char*) file_iterator, 
				 length - file_size_ctr, virname, sig_hashtable))
	{
	  return CL_VIRUS;
	}
      }
    }

    file_size_ctr++;
    file_iterator++;
  }

  return 0;
}
